package cn.fetosoft.woodpecker.core.jmeter.processor;

import cn.fetosoft.commons.secret.HmacSHA1Signer;
import cn.fetosoft.commons.utils.MD5Util;
import org.apache.commons.lang3.StringUtils;
import org.apache.jmeter.config.Arguments;
import org.apache.jmeter.extractor.json.jsonpath.JSONManager;
import org.apache.jmeter.processor.PreProcessor;
import org.apache.jmeter.protocol.http.sampler.HTTPSamplerBase;
import org.apache.jmeter.samplers.Sampler;
import org.apache.jmeter.testelement.AbstractTestElement;
import org.apache.jmeter.testelement.property.CollectionProperty;
import org.apache.jmeter.testelement.property.JMeterProperty;
import org.apache.jmeter.testelement.property.PropertyIterator;
import org.apache.jmeter.threads.JMeterContext;
import org.apache.jmeter.threads.JMeterVariables;
import org.apache.jorphan.util.JOrphanUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.CollectionUtils;

import java.io.Serializable;
import java.text.ParseException;
import java.util.*;

/**
 * 参数签名处理器
 * @author guobingbing
 * @wechat t_gbinb
 * @since 2021/8/8 15:52
 */
public class ParamsSignProcessor extends AbstractTestElement implements PreProcessor, Serializable {

    private static final Logger log = LoggerFactory.getLogger(ParamsSignProcessor.class);
    private static final String SEPARATOR = ";";
    private static final String JOINER = "&";
    private static final String SIGN_MD5 = "md5";
    private static final String SIGN_HMAC = "hmac";
    private static final ThreadLocal<JSONManager> localMatcher = ThreadLocal.withInitial(JSONManager::new);
    private static final String SIGN_TYPE = "ParamsSignProcessor.signType";
    private static final String SIGN_KEY = "ParamsSignProcessor.signKey";
    private static final String SIGN_NAME = "ParamsSignProcessor.signName";
    private static final String SIGN_JOINER = "ParamsSignProcessor.signJoiner";
    private static final String PARAM_NAME = "ParamsSignProcessor.paramName";
    private static final String JSON_PATH = "ParamsSignProcessor.jsonPath";

    public String getSignType() {
        return getPropertyAsString(SIGN_TYPE);
    }

    public void setSignType(String signType) {
        setProperty(SIGN_TYPE, signType);
    }

    public String getSignKey() {
        return getPropertyAsString(SIGN_KEY);
    }

    public void setSignKey(String signKey) {
        setProperty(SIGN_KEY, signKey);
    }

    public String getSignName() {
        return getPropertyAsString(SIGN_NAME);
    }

    public void setSignName(String signName) {
        setProperty(SIGN_NAME, signName);
    }

    public String getSignJoiner() {
        String signJoiner = getPropertyAsString(SIGN_JOINER, JOINER);
        return StringUtils.isBlank(signJoiner)?JOINER:signJoiner;
    }

    public void setSignJoiner(String signJoiner) {
        setProperty(SIGN_JOINER, signJoiner);
    }

    public String getParamName() {
        return getPropertyAsString(PARAM_NAME);
    }

    public void setParamName(String paramName) {
        setProperty(PARAM_NAME, paramName);
    }

    public String getJsonPath() {
        return getPropertyAsString(JSON_PATH);
    }

    public void setJsonPath(String jsonExpression) {
        setProperty(JSON_PATH, jsonExpression);
    }

    @Override
    public void process() {
        JMeterContext context = getThreadContext();
        Sampler sam = context.getCurrentSampler();
        HTTPSamplerBase sampler;
        if(!(sam instanceof HTTPSamplerBase)) {
            log.error("Can't apply HTML Link Parser when the previous sampler run is not an HTTP Request.");
            return;
        }else {
            sampler = (HTTPSamplerBase) sam;
        }

        String[] paramNames = this.getParamName().split(SEPARATOR);
		if(paramNames.length<1){
			log.error("参数个数不能为空！");
			return;
		}
		Arguments args = sampler.getArguments();
		if(args.getArguments().size()<1){
			log.error("POST内容不能为空！");
			return;
		}
		try{
			String signData = "";
			if(sampler.getPostBodyRaw()){
				String[] jsonPath = this.getJsonPath().split(SEPARATOR);
				if(paramNames.length!=jsonPath.length){
					log.error("参数个数与JSON提取表达式个数不匹配！");
					return;
				}
				signData = this.jsonToSign(args, paramNames, jsonPath);
			}else{
				signData = this.formToSign(args, paramNames, sampler);
			}
			JMeterVariables vars = getThreadContext().getVariables();
			vars.put(getSignName(), signData);
		}catch(Exception e){
			log.error("ParamsSignProcessor", e);
		}
    }

	/**
	 * 从json中提取参数
	 * @param args
	 * @param signParams
	 * @param expressions
	 * @return
	 */
    private String jsonToSign(Arguments args, String[] signParams,
							String[] expressions) throws ParseException{
    	String json = args.getArgument(0).getValue();
        Map<String, String> map = new HashMap<>(32);
		String name = "", jsonPath = "";
		for(int i=0; i<signParams.length; i++) {
			name = signParams[i];
			jsonPath = expressions[i];
			List<Object> values = localMatcher.get()
					.extractWithJsonPath(json, jsonPath);
			if(CollectionUtils.isEmpty(values)){
				throw new ParseException(jsonPath, i);
			}
			map.put(name, values.get(0).toString());
		}
		String signData = this.sign(map);
		String search = "${" + getSignName() + "}";
		String arguments = JOrphanUtils.replaceFirst(json, search, signData);
		//arguments = StringUtilities.substitute(arguments, search, signData);
		args.getArgument(0).setValue(arguments);
		return signData;
    }

	/**
	 * 将form参数转换为map
	 * @param args
	 * @param signParams
	 * @return
	 * @throws ParseException
	 */
    private String formToSign(Arguments args, String[] signParams,
							HTTPSamplerBase sampler) throws ParseException{
		CollectionProperty props = args.getArguments();
		PropertyIterator iter = props.iterator();
		Map<String, String> propMap = new HashMap<>();
		while (iter.hasNext()){
			JMeterProperty p = iter.next();
			String value = p.getStringValue();
			String[] kv = value.split("=");
			propMap.put(kv[0], kv[1]);
		}
		Map<String, String> map = new HashMap<>(32);
		for(String name : signParams){
			map.put(name, propMap.get(name));
		}
		String signData = this.sign(map);
		args.getArguments().remove(getSignName());
		sampler.addEncodedArgument(getSignName(), signData);
		return signData;
	}

	/**
	 * 数据签名
	 * @param map
	 * @return
	 */
	private String sign(Map<String, String> map){
		String[] sortedKeys = map.keySet().toArray(new String[]{});
		Arrays.sort(sortedKeys);
		StringJoiner joiner = new StringJoiner(getSignJoiner());
		for(String key : sortedKeys){
			joiner.add(key + "=" + map.get(key));
		}
		String signData = joiner.toString();
		if(SIGN_HMAC.equalsIgnoreCase(getSignType())){
			signData = HmacSHA1Signer.signString(signData, getSignKey());
		}else{
			signData = MD5Util.md5(signData);
		}
		return signData;
	}
}
